package util

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"github.com/ethereum/go-ethereum/crypto/ecies"
	mathRand "math/rand"
	"os"
	"strings"
	"time"
)

// 生成指定math/rand字节长度的随机字符串
func GetRandomString(length int) string {
	str := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ~!@#$%^&*()_+?=-"
	bytes := []byte(str)
	result := []byte{}

	r := mathRand.New(mathRand.NewSource(time.Now().UnixNano()))
	for i := 0; i < length; i++ {
		result = append(result, bytes[r.Intn(len(bytes))])
	}
	return string(result)
}

// 生成椭圆曲线非对称密钥对
func GenerateECDSAKey(randKey string) {
	//根据随机字符串长度设置curve曲线
	length := len(randKey)
	//elliptic包实现了几条覆盖素数有限域的标准椭圆曲线,Curve代表一个短格式的Weierstrass椭圆曲线，其中a=-3
	if length < 224/8 {
		err := errors.New("私钥长度太短，至少为36位！")
		panic(err)
	}
	var curve elliptic.Curve
	if length >= 521/8+8 {
		//长度大于73字节，返回一个实现了P-512的曲线
		curve = elliptic.P521()
	} else if length >= 384/8+8 {
		//长度大于56字节，返回一个实现了P-384的曲线
		curve = elliptic.P384()
	} else if length >= 256/8+8 {
		//长度大于40字节，返回一个实现了P-256的曲线
		curve = elliptic.P256()
	} else if length >= 224/8+8 {
		//长度大于36字节，返回一个实现了P-224的曲线
		curve = elliptic.P224()
	}
	//1.私钥序列化到磁盘
	privateKey, err := ecdsa.GenerateKey(curve, strings.NewReader(randKey))
	if err != nil {
		panic(err)
	}
	//x509编码
	x509PriBytes, err := x509.MarshalECPrivateKey(privateKey)
	if err != nil {
		panic(err)
	}
	block := pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: x509PriBytes,
	}
	file, err := os.Create("ec-private.pem")
	if err != nil {
		panic(err)
	}
	pem.Encode(file, &block)
	file.Close()
	fmt.Println("ec-private.pem is Create")
	//2.公钥序列化到磁盘
	publicKey := privateKey.PublicKey
	x509PubBytes, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		panic(err)
	}
	block = pem.Block{
		Type:  "EC PUBLIC KEY",
		Bytes: x509PubBytes,
	}
	file, err = os.Create("ec-public.pem")
	if err != nil {
		panic(err)
	}
	pem.Encode(file, &block)
	file.Close()
	fmt.Println("ec-public.pem is Create")
}

// ECDSA 公钥从文件中提取
func ECDSAFromFileGetPublicKey(fileName string) *ecdsa.PublicKey {
	//将公钥文件中的公钥读出, 得到使用pem编码的字符串-- 读文件
	file, err := os.Open(fileName)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	fileInfo, err := file.Stat()
	if err != nil {
		panic(err)
	}
	bytes := make([]byte, fileInfo.Size())
	_, err = file.Read(bytes)
	if err != nil {
		panic(err)
	}
	//将得到的字符串解码-- pem.Decode
	decodePubKey, _ := pem.Decode(bytes)
	//使用x509将编码之后的公钥解析出来
	publicKey, err := x509.ParsePKIXPublicKey(decodePubKey.Bytes)
	if err != nil {
		panic(err)
	}
	//断言转换类型
	key := publicKey.(*ecdsa.PublicKey)
	return key
}

// ECDSA 私钥从文件中提取
func ECDSAFromFileGetPrivateKey(fileName string) *ecdsa.PrivateKey {
	//将私钥文件中的私钥读出, 得到使用pem编码的字符串-- 读文件
	file, err := os.Open(fileName)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	fileInfo, err := file.Stat()
	if err != nil {
		panic(err)
	}
	bytes := make([]byte, fileInfo.Size())
	_, err = file.Read(bytes)
	if err != nil {
		panic(err)
	}
	//将得到的字符串解码-- pem.Decode
	decodePriKey, _ := pem.Decode(bytes)
	//使用x509将编码之后的公钥解析出来
	privateKey, err := x509.ParseECPrivateKey(decodePriKey.Bytes)
	if err != nil {
		panic(err)
	}
	return privateKey
}

// ECCEncrypt 椭圆曲线公钥加密数据
func ECCEncrypt(plainText []byte, fileName string) ([]byte, error) {
	publicKey := ECDSAFromFileGetPublicKey(fileName)
	//读取文件的ecdsa公钥转化成ecies公钥
	publicKeyForEcies := ecies.ImportECDSAPublic(publicKey)
	return ecies.Encrypt(rand.Reader, publicKeyForEcies, plainText, nil, nil)
}

// ECCDecrypt 椭圆曲线私钥解密数据
func ECCDecrypt(cipher []byte, fileName string) ([]byte, error) {
	privateKey := ECDSAFromFileGetPrivateKey(fileName)
	//读取文件的ecdsa私钥转化成ecies私钥
	privateKeyForEcies := ecies.ImportECDSA(privateKey)
	if src, err := privateKeyForEcies.Decrypt(cipher, nil, nil); err != nil {
		return src, err
	} else {
		return src, nil
	}
}
